import os
import json
import requests
from volcenginesdkarkruntime import Ark

class VideoGenerator:
    def __init__(self, api_key, base_url="https://ark.cn-beijing.volces.com/api/v3"):
        self.client = Ark(
            base_url=base_url,
            api_key=api_key,
        )
        self.base_url = base_url
    
    def generate_video(self, prompt):
        try:
            # 调用内容生成API创建视频生成任务
            response = self.client.content_generation.tasks.create(
                model="doubao-seedance-1-0-pro-250528",
                content=[{"text": prompt, "type": "text"}]
            )
            print(f"返回的响应数据为{response}")
            # 返回任务ID
            return {
                "task_id": response.id
            }
        except Exception as e:
            raise Exception(f"视频生成失败: {str(e)}")
    
    def get_video_generation_status(self, task_id):
        try:
            # 使用requests库直接调用API查询任务状态
            url = f"{self.base_url}/contents/generations/tasks/{task_id}"
            headers = {
                "Authorization": f"Bearer {self.client.api_key}"
            }
            response = requests.get(url, headers=headers)
            response.raise_for_status()  # 如果响应状态码不是200，会抛出异常
            
            # 解析响应内容
            response_data = response.json()
            
            # 返回任务状态和完整响应内容
            return {
                "task_id": response_data.get("id", task_id),
                "status": response_data.get("status"),
                "result": response_data.get("result"),
                "content": response_data.get("content"),
                "full_response": response_data
            }
        except Exception as e:
            raise Exception(f"获取视频生成状态失败: {str(e)}")
    
    def generate_video_from_image(self, prompt, image_url):
        try:
            # 调用内容生成API创建图生视频任务
            response = self.client.content_generation.tasks.create(
                model="doubao-seedance-1-0-pro-250528",
                content=[{"text": prompt, "type": "text"}, {"image_url": {"url": image_url}, "type": "image_url"}]
            )
            
            # 返回任务ID
            return {
                "task_id": response.id
            }
        except Exception as e:
            raise Exception(f"图生视频生成失败: {str(e)}")
